#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
import re


# 连接MongoDB
def InitMongoDB(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME):
    client = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = client[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return client, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


# 第一步导入数据后清除空字段
def Clean_eLMSG():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180212")

    for result in collection_eLMSG.find():
        del_empty_feild(result)
        collection_eLMSG.replace_one({"_id": result["_id"]}, result)

    CloseMongoDB(client_eLMSG)


def del_empty_feild(_dict: dict, parent=None, key_in_parent=None):
    for k in list(_dict.keys()):
        v = _dict[k]
        if isinstance(v, str):
            if v is None or not len(v.strip()):
                del _dict[k]
        elif isinstance(v, list):
            if not len(v):
                del _dict[k]
            else:
                for i in v:
                    if isinstance(i, dict):
                        del_empty_feild(i)
        elif isinstance(v, dict):
            del_empty_feild(v, _dict, k)
        elif v is None:
            del _dict[k]

    if not len(_dict) and parent is not None and key_in_parent is not None:
        del parent[key_in_parent]


# 第二步补全Taxonomy信息
def Fullfill_Taxonomy():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180129")
    client_Taxonomy, collection_Taxonomy = InitMongoDB("10.188.188.22", 27017, "Dictionary", "taxonomy_20180123")

    for result in collection_eLMSG.find({"TaxonID": {"$exists": 1}}):
        InTaxonomy = collection_Taxonomy.find_one({"tax_id": result["TaxonID"]})
        if InTaxonomy != None:
            result2 = {}
            for key in InTaxonomy.keys():
                if key.startswith("names_"):
                    result2[key] = InTaxonomy[key]
                if key.startswith("division_"):
                    result2[key] = InTaxonomy[key]
                if key.startswith("genetic_"):
                    result2[key] = InTaxonomy[key]
                if key == "rank":
                    result2["Rank"] = InTaxonomy["rank"]
                if key == "parent_tax_id":
                    result2["parent_tax_id"] = InTaxonomy["parent_tax_id"]["id"]
            collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": result2})
        else:
            # print(result["TaxonID"])
            OldTaxonomy = collection_Taxonomy.find_one({"old_tax_id": result["TaxonID"]})
            result3 = {}
            for key in OldTaxonomy.keys():
                if key.startswith("names_"):
                    result3[key] = OldTaxonomy[key]
                if key.startswith("division_"):
                    result3[key] = OldTaxonomy[key]
                if key.startswith("genetic_"):
                    result3[key] = OldTaxonomy[key]
                if key == "rank":
                    result3["Rank"] = OldTaxonomy["rank"]
                if key == "tax_id":
                    result3["TaxonID"] = OldTaxonomy["tax_id"]
                if key == "parent_tax_id":
                    result3["parent_tax_id"] = OldTaxonomy["parent_tax_id"]["id"]
            collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": result3})

    CloseMongoDB(client_eLMSG)
    CloseMongoDB(client_Taxonomy)


# 第三步从Taxonomy补齐其他Rank
def ImportOtherRanksFromTaxonomy():
    # MicrobiomeRoots = set(["2157", "2", "4751", "10239", "12884", "151659", "408169","1"])
    MicrobiomeRoots = set(["1"])

    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180129")
    client_Taxonomy, collection_Taxonomy = InitMongoDB("10.188.188.22", 27017, "Dictionary", "taxonomy_20180123")

    for result in collection_eLMSG.find({"TaxonID": {"$exists": 1}}):
        RecursionTaxon(result["TaxonID"], collection_Taxonomy, collection_eLMSG, MicrobiomeRoots)

    CloseMongoDB(client_eLMSG)
    CloseMongoDB(client_Taxonomy)


def RecursionTaxon(TaxonID, collection_Taxonomy, collection_eLMSG, MicrobiomeRoots):
    parentTaxonID = collection_Taxonomy.find_one({"dms_id": TaxonID})["parent_tax_id"]["id"]

    if collection_eLMSG.find_one({"TaxonID": parentTaxonID}) == None:
        InTaxonomy = collection_Taxonomy.find_one({"dms_id": parentTaxonID})
        result = {}
        for key in InTaxonomy.keys():
            if key.startswith("names_"):
                result[key] = InTaxonomy[key]
            if key.startswith("division_"):
                result[key] = InTaxonomy[key]
            if key.startswith("genetic_"):
                result[key] = InTaxonomy[key]
            if key == "rank":
                result["Rank"] = InTaxonomy["rank"]
            if key == "tax_id":
                result["TaxonID"] = InTaxonomy["tax_id"]
            if key == "parent_tax_id":
                result["parent_tax_id"] = InTaxonomy["parent_tax_id"]["id"]
        collection_eLMSG.insert_one(result)
    if parentTaxonID not in MicrobiomeRoots:
        RecursionTaxon(parentTaxonID, collection_Taxonomy, collection_eLMSG, MicrobiomeRoots)


def GenerateStructuredReference():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180212")

    for result in collection_eLMSG.find({"References": {"$exists": 1}}, {"References": 1}):
        reference = []
        for line in result["References"]:
            reference.append(GetReference(line))
        collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"StructuredReference": reference}})
    CloseMongoDB(client_eLMSG)


def GetReference(line):
    import re
    reference = {}
    regex = "\(((PM|DO).*)\)"
    prog = re.compile(regex)
    result = prog.search(line)
    if result:
        ids = result.group(1)
        for id in re.split(r'[\s ]', ids):
            if id.startswith("PMID"):
                reference["PMID"] = id[5:]
            elif id.startswith("DOI"):
                doi = id[4:]
                doi = re.sub("<.*?>", "", doi)
                reference["DOI"] = doi
            elif id.startswith("PMCID"):
                reference["PMCID"] = id[6:]
        txt = line.replace("(" + ids + ")", "", -1).strip()
        if txt != "":
            reference["txt"] = txt
    else:
        txt = line.strip()
        if txt != "":
            reference["txt"] = txt
    return reference


# 补全LMSG_ID
def FullfillLMSG_ID(startID):
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180212")

    for result in collection_eLMSG.find({"LMSG_ID": {"$exists": 0}}, {"_id": 1}):
        startID += 1
        lmsg_id = "MSG" + str(startID).zfill(6)
        collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"LMSG_ID": lmsg_id}})

    CloseMongoDB(client_eLMSG)


# 测试Lineage与Name对应关系
def CheckLineage():
    ht = set()

    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180212")
    for result in collection_eLMSG.find({}, {"_id": 0, "Lineage": 1}):
        for taxon in result["Lineage"].split("."):
            ht.add(taxon)

    for taxon in ht:
        result = collection_eLMSG.find_one({"Name": taxon})
        if result == None:
            print(taxon)
    CloseMongoDB(client_eLMSG)


def CheckUniqueName():
    ht = set()
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180212")
    for result in collection_eLMSG.find({}, {"Name": 1, "_id": 0}):
        if result["Name"] in ht:
            print(result["Name"])
        else:
            ht.add(result["Name"])
    CloseMongoDB(client_eLMSG)


def FullfillParentLMSGId():
    nameset = set(["unC", "unO", "unF", "unP", "UnC", "UnO", "UnF", "UnP"])
    ht = {}
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180212")
    for result in collection_eLMSG.find({}, {"Name": 1, "LMSG_ID": "1", "_id": 0}):
        ht[result["Name"]] = result["LMSG_ID"]
    for result in collection_eLMSG.find({"Parent_LMSG_ID": {"$exists": 0}}, {"Lineage": 1}):
        lineage = result["Lineage"].split(".")
        for i in range(len(lineage) - 1, -1, -1):
            if lineage[i] in nameset:
                lineage.pop(i)
        if len(lineage) > 2:
            collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"Parent_LMSG_ID": ht[lineage[-2]]}})
    CloseMongoDB(client_eLMSG)


def Structure_Genbank_16SrRNA():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180212")
    for result in collection_eLMSG.find({"16SrRNA": {"$exists": 1}}, {"16SrRNA": 1}):
        collection_eLMSG.update_one({"_id": result["_id"]},
                                    {"$set": {"Genbank_16SrRNA": result["16SrRNA"]}, "$unset": {"16SrRNA": 1}})
    CloseMongoDB(client_eLMSG)


def Structure_JGI_IMG():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180212")
    for result in collection_eLMSG.find({"JGI IMG": {"$exists": 1}}, {"JGI IMG": 1}):
        collection_eLMSG.update_one({"_id": result["_id"]},
                                    {"$set": {"JGI_IMG": result["JGI IMG"]}, "$unset": {"JGI IMG": 1}})
    CloseMongoDB(client_eLMSG)


def Structure_NCBI_Assembly():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180212")
    for result in collection_eLMSG.find({"NCBI GenBank": {"$exists": 1}}, {"NCBI GenBank": 1}):
        collection_eLMSG.update_one({"_id": result["_id"]},
                                    {"$set": {"NCBI_Assembly": result["NCBI GenBank"]}, "$unset": {"NCBI GenBank": 1}})
    CloseMongoDB(client_eLMSG)


def ChangeTypeFieldType():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180212")

    for result in collection_eLMSG.find({"Type": {"$exists": 1}}, {"Type": 1}):
        if isinstance(result["Type"], str):
            type = result["Type"]
            id = result["_id"]
            collection_eLMSG.update_one({"_id": id}, {"$set": {"Type": [type]}})

    CloseMongoDB(client_eLMSG)


def StructureFattyAcids():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180212")
    for result in collection_eLMSG.find({"Phenotype.Biochemistry.fattyacids": {"$exists": 1}},
                                        {"Phenotype.Biochemistry.fattyacids": 1}):
        fattyacids = trans_fattyacids(result["Phenotype"]["Biochemistry"]["fattyacids"])
        collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"Phenotype.Biochemistry.fattyacids": fattyacids}})
    CloseMongoDB(client_eLMSG)


def trans_fattyacids(fattyacids_str):
    import json
    fattyacids = []
    for k, v in json.loads(fattyacids_str).items():
        fattyacids.append({"name": k, "account": str(v)})
    return fattyacids


def AddFunctionAnnotation(filepath):
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180212")

    infile = open(filepath, "r")
    infile.readline()
    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        result = collection_eLMSG.find_one({"NCBI_Assembly.id": splitline[0]}, {"_id": 1})
        if result != None:
            gcaId = splitline[0].replace(".", "_")
            anno = {"Uniprot": int(splitline[4]), "TIGRFAM": int(splitline[5]), "KEGG": int(splitline[2]),
                    "COGs": int(splitline[1]),
                    "Pfam": int(splitline[3])}
            collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"FunctionAnnotation." + gcaId: anno}})
    infile.close()
    CloseMongoDB(client_eLMSG)


if __name__ == '__main__':
    Clean_eLMSG()
    # CheckLineage()
    # CheckUniqueName()
    # FullfillLMSG_ID(0)
    # FullfillParentLMSGId()
    # GenerateStructuredReference()
    # Structure_Genbank_16SrRNA()
    # Structure_JGI_IMG()
    # Structure_NCBI_Assembly()
    # ChangeTypeFieldType()
    # StructureFattyAcids()
    # AddFunctionAnnotation("D:/anno_stat.txt")
